package cas.ihep.hep.analysis;

import cas.ihep.hep.util.MyApplication;
import cas.ihep.spark.context.SparkContextExt;
import cas.ihep.spark.util.PartitionLocator;
import cas.ihep.util.Finalizer;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.VoidFunction;
import sun.misc.Unsafe;
import sun.nio.ch.DirectBuffer;

import java.io.IOException;
import java.io.Serializable;
import java.util.*;

public class TestDemo extends MyApplication implements Serializable {
    public TestDemo(String[] args) {
        super(args);
    }

    @Override
    public void run() throws Exception {
        srun();
    }

    static class Data implements Serializable{
        String host;
        Data(String str){
            host=str;
        }

        public String toString(){
            return host;
        }
    }

    private static final String[] hosts = new String[]{
            "besh01.ihep.ac.cn", "besh02.ihep.ac.cn", "besh03.ihep.ac.cn", "besh04.ihep.ac.cn",
            "besh05.ihep.ac.cn", "besh06.ihep.ac.cn", "besh07.ihep.ac.cn", "besh08.ihep.ac.cn"
    };

    private static final HashMap<String,String> host2Ip=new HashMap<>();

    static {
        host2Ip.put("besh01.ihep.ac.cn","192.168.60.108");
        host2Ip.put("besh01","192.168.60.108");
        host2Ip.put("besh02.ihep.ac.cn","192.168.60.109");
        host2Ip.put("besh02","192.168.60.109");
        host2Ip.put("besh03.ihep.ac.cn","192.168.60.110");
        host2Ip.put("besh03","192.168.60.110");
        host2Ip.put("besh04.ihep.ac.cn","192.168.60.115");
        host2Ip.put("besh04","192.168.60.115");
        host2Ip.put("besh05.ihep.ac.cn","192.168.60.118");
        host2Ip.put("besh05","192.168.60.118");
        host2Ip.put("besh06.ihep.ac.cn","192.168.60.119");
        host2Ip.put("besh06","192.168.60.119");
        host2Ip.put("besh07.ihep.ac.cn","192.168.60.120");
        host2Ip.put("besh07","192.168.60.120");
        host2Ip.put("besh08.ihep.ac.cn","192.168.60.121");
        host2Ip.put("besh08","192.168.60.121");
    }

    private static void srun()throws IOException {
        try(Finalizer closer=Finalizer.create()){
            Unsafe unsafe;
            DirectBuffer db;
            SparkContextExt jsc=closer.register(new SparkContextExt(new SparkConf()));
            ArrayList<Data> data=new ArrayList<>();
            Random rand=new Random();
            for(int i=0;i<32;i++){
                data.add(new Data(host2Ip.get(hosts[rand.nextInt(hosts.length)])));
            }
            JavaRDD<Data> dataRdd=jsc.wholeList(data, new PartitionLocator<Data>() {
                @Override
                public List<String> getPreferredLocations(List<Data> list) {
                    int sz=list.size();
                    if(sz==0){
                        //noinspection unchecked
                        return Collections.EMPTY_LIST;
                    }else if(sz==1){
                        return Collections.singletonList(list.get(0).toString());
                    }
                    List<String> ret=new ArrayList<>(sz);
                    for(Data dt:list){
                        ret.add(dt.toString());
                    }
                    return ret;
                }
            });

            dataRdd.foreach(new VoidFunction<Data>() {
                @Override
                public void call(Data data) {
                    System.out.println(data);
                }
            });
        }
    }
}
